# Copyright (c) Microsoft. All rights reserved.
# Licensed under the MIT license. See LICENSE file in the project root for full license information.

"""
Version 2.2
what's new:
- add GetLatestAverageLoss() function
"""

import numpy as np
import time
import math
import os
import sys

from .ActivatorFunction_2_0 import *
from .ClassifierFunction_2_0 import *
from .EnumDef_2_0 import *
from .TrainingHistory_2_2 import *
from .LossFunction_1_1 import *
from .WeightsBias_1_0 import *
from .WeightsBias_1_0 import *
# import EnumDef_2_0


# 有些重复的容易理解的就没写注释
class NeuralNet_2_2(object):
    def __init__(self, hp, model_name):
        """
        初始化神经网络
        :param hp: 神经网络各参数
        :param model_name: 模型名称（已经训练得到的神经网络权重和偏置）
        """
        self.hp = hp  # 神经网络各参数
        self.model_name = model_name  # 模型文件名称，用于读取模型和保存模型
        # 设置模型路径，用于读取初始化的权重和偏置、保存模型（训练后得到的神经网络各层权重和偏置）
        self.subfolder = os.getcwd() + "\\" + self.__create_subfolder()
        print(self.subfolder)

        # 初始化第一层的权重和偏置
        self.wb1 = WeightsBias_1_0(self.hp.num_input, self.hp.num_hidden, self.hp.init_method, self.hp.eta)
        self.wb1.InitializeWeights(self.subfolder, False)
        # 初始化第二层的权重和偏置
        self.wb2 = WeightsBias_1_0(self.hp.num_hidden, self.hp.num_output, self.hp.init_method, self.hp.eta)
        self.wb2.InitializeWeights(self.subfolder, False)

    def __create_subfolder(self):
        if self.model_name != None:
            path = self.model_name.strip()
            path = path.rstrip("\\")
            isExists = os.path.exists(path)
            if not isExists:
                os.makedirs(path)
            return path

    def forward(self, batch_x):
        """
        前向计算，从输入层到隐藏层到输出层
        :param batch_x: 输入，批次的训练数据（也可以是验证集、预测集）
        :return:
        """
        # 计算隐藏层输出
        self.Z1 = np.dot(batch_x, self.wb1.W) + self.wb1.B
        # 隐藏层的输出经过激活函数，作为输出层的输入
        self.A1 = Sigmoid().forward(self.Z1)
        # 计算输出层的输出
        self.Z2 = np.dot(self.A1, self.wb2.W) + self.wb2.B
        # 输出层的输出经过激活函数
        # 如果想要的是二分类结果就用BinaryClassifier，激活函数就是Logistic
        if self.hp.net_type == NetType.BinaryClassifier:
            self.A2 = Logistic().forward(self.Z2)
        # 如果想要的是多重分类就用MultipleClassifier，激活函数就是SoftMax
        elif self.hp.net_type == NetType.MultipleClassifier:
            self.A2 = SoftMax().forward(self.Z2)
        else:
        # 如果是简单线性回归任务，就无需激活函数
            self.A2 = self.Z2
        self.output = self.A2


    def backward(self, batch_x, batch_y, batch_a):
        """
        反向计算，通过真实值、预测值（第二层输出）和第一层输入值计算每层的权重和偏置的梯度变化
        :param batch_x: 第一层输入
        :param batch_y: 真实值
        :return:
        """
        m = batch_x.shape[0]
        dZ2 = batch_a - batch_y  # 误差
        # 计算输出层权重的梯度
        self.wb2.dW = np.dot(self.A1.T, dZ2) / m
        # 计算输出层偏置的梯度
        self.wb2.dB = np.sum(dZ2, axis=0, keepdims=True) / m
        # 计算中间梯度
        d1 = np.dot(dZ2, self.wb2.W.T)
        # 反向传播经过激活函数，得到输出层输出的梯度变化
        dZ1, _ = Sigmoid().backward(None, self.A1, d1)
        # 计算隐藏层dW的梯度
        self.wb1.dW = np.dot(batch_x.T, dZ1) / m
        # 计算隐藏层dB的梯度
        self.wb1.dB = np.sum(dZ1, axis=0, keepdims=True)

    def updata(self):
        """
        更新权重和偏置
        :return:
        """
        self.wb1.Update()
        self.wb2.Update()

    def inference(self, x):
        """
        模型预测函数
        :param x:
        :return:
        """
        self.forward(x)
        return self.output

    def train(self, dataReader, checkpoint, need_test):
        """
        读取数据，前向计算，反向传播，更新权重和偏置，计算训练集和验证集在当前检查点的损失和准确率（低于eps停止训练）
        :param dataReader: 数据读取对象
        :param checkpoint: 检查点，每到一个检查点就计算损失和准确率
        :param need_test: 是否需要测试集
        :return:
        """
        # 创建保存训练集和验证集的历史损失值和准确率的对象
        self.loss_trace = TrainingHistory_2_2()
        # 初始化损失函数类型
        self.loss_func = LossFunction_1_1(self.hp.net_type)
        if self.hp.batch_size == -1:
            self.hp.batch_size = dataReader.num_train
        max_iteration = math.ceil(dataReader.num_train / self.hp.batch_size)  # 每一轮的最大训练次数
        # 将定点轮次检查转换成定点迭代次数检查
        checkpoint_iteration = int(max_iteration * checkpoint)
        need_stop = False  # 是否需要停止训练变量
        for epoch in range(self.hp.max_epoch):
            dataReader.Shuffle()  # 打乱所有数据集
            # 进行每一轮训练
            for iteration in range(max_iteration):
                # 进行每一轮中每一次训练
                batch_x, batch_y = dataReader.GetBatchTrainSamples(self.hp.batch_size, iteration)  # 获取该批次的训练集和真实值集合
                batch_a = self.inference(batch_x)  # 获取本轮前向计算后第二层的输出
                self.backward(batch_x, batch_y, batch_a)  # 反向传播
                self.updata()  # 更新两层的权重和偏置
                total_iteration = epoch * max_iteration + iteration  # 当前迭代次数，目前的训练次数总和
                # 如果当前迭代正是需要检查的迭代
                if (total_iteration+1) % checkpoint_iteration == 0:
                    # 计算并记录当前模型对训练集和验证集的损失和准确率，并且判断损失是否达到eps要求
                    need_stop = self.CheckErrorAndLoss(dataReader, batch_x, batch_y, epoch, total_iteration)
                    if need_stop:
                        break
            if need_stop:
                break
        self.SaveResult()  # 保存结果，经过训练完后的每一层的权重和偏置
        if need_test:  # 如果需要测试，那么就计算该模型的整体准确率
            print("testing......")
            accuracy = self.Test(dataReader)
            print(f"accuracy={accuracy}")


    def CheckErrorAndLoss(self, dataReader, train_x, train_y, epoch, total_iteration):
        """
        检查此时模型的训练集和验证集的损失值和准确率并返回是否需要停止训练
        :param dataReader: 数据读取对象
        :param train_x: 批次输入数据
        :param train_y: 批次真实值
        :param epoch: 本次训练轮次
        :param total_iteration: 本次训练总迭代次数
        :return: 是否需要停止训练
        """
        print(f"epoch={epoch}, total_iteration={total_iteration}")
        # 对于当前批次输入的模型预测
        self.forward(train_x)
        # 这里是使用多重分类，所以使用的交叉熵损失函数计算损失
        loss_train = self.loss_func.CheckLoss(self.output, train_y)
        # 计算训练集准确率
        accuracy_train = self.__CalAccuracy(self.output, train_y)
        print(f"loss_train={loss_train}, accuracy_train={accuracy_train}")
        # 读取验证集数据
        vld_x, vld_y = dataReader.GetValidationSet()
        # 对于验证集的模型预测
        self.forward(vld_x)
        # 同上，使用交叉熵损失函数计算验证集的损失
        loss_vld = self.loss_func.CheckLoss(self.output, vld_y)
        # 计算验证集的准确率
        accuracy_vld = self.__CalAccuracy(self.output, vld_y)
        print(f"loss_valid={loss_vld}, accuracy_valid={accuracy_vld}")
        # 记录训练集和验证集的损失值和准确率
        need_stop = self.loss_trace.Add(epoch, total_iteration, loss_train, accuracy_train, loss_vld, accuracy_vld)
        # 判断模型此时是否低于等于我们的损失要求
        if loss_vld <= self.hp.eps:
            need_stop = True
        return need_stop

    def Test(self, dataReader):
        """
        从样本中抽出部分作为测试集来查看当前模型的精确度
        :param dataReader: 数据读取对象
        :return: 当前模型测试的正确率
        """
        x, y = dataReader.GetTestSet()  # 获取测试数据集
        self.forward(x)  # 对于测试机的模型预测
        correct = self.__CalAccuracy(self.output, y)  # 计算对于测试机的zhunquel
        print(f"correct={correct}")
        return correct


    def __CalAccuracy(self, a, y):
        """
        根据模型输出（预测值）和真实值计算模型的准确率
        :param a: 模型的输出（预测值）
        :param y: 真实值
        :return: 当前模型的准确率
        """
        assert a.shape == y.shape  # 断言，如果预测值和真实值的形状不同，则产生错误
        m = a.shape[0]
        # 简单线性拟合，计算评估指标，这里越接近1越好
        if self.hp.net_type == NetType.Fitting:
            var = np.var(y)
            mse = np.sum((a-y)**2) / m
            r2 = 1 - mse / var
            return r2
        # 如果是二分类情况，计算正确预测的个数，也就是准确率
        elif self.hp.net_type == NetType.BinaryClassifier:
            b = np.round(a)
            r = (b == y)
            correct = r.sum()
            return correct / m
        # 如果是多重分类的情况，也是计算正确预测的个数，巨额救市准确率
        elif self.hp.net_type == NetType.MultipleClassifier:
            ra = np.argmax(a, axis=1)
            ry = np.argmax(y, axis=1)
            r = (ra == ry)
            correct = r.sum()
            return correct / m

    def SaveResult(self):
        """
        保存当前模型（每层的权重和偏置）
        :return:
        """
        self.wb1.SaveResultValue(self.subfolder, "wb1")
        self.wb2.SaveResultValue(self.subfolder, "wb2")

    def LoadResult(self):
        """
        加载模型，读取模型中每层的权重和偏置
        :return:
        """
        self.wb1.LoadResultValue(self.subfolder, "wb1")
        self.wb2.LoadResultValue(self.subfolder, "wb2")

    def ShowTrainingHistory(self):
        """
        可视化训练过程中损失和准确率变化
        :return:
        """
        self.loss_trace.ShowLossHistory(self.hp)

    def GetTrainingHistory(self):
        """
        获取训练的损失和准确率历史记录实例对象
        :return: 训练的损失和准确率历史记录实例对象
        """
        return self.loss_trace

    def GetEpochNumber(self):
        """
        获取当前训练的轮次
        :return: 当前训练的轮次
        """
        return self.loss_trace.GetEpochNumber()

    def GetLatestAverageLoss(self, count=10):
        """
        获取最后count次的验证集平均损失（count>保存的个数则返回最后一次损失值）
        :param count: 获取最后count次的验证集平均损失
        :return: 最后count次的验证集平均损失（count>保存的个数则返回最后一次损失值）
        """
        return self.loss_trace.GetLatestAverageLoss(count)



